import d2lzh as d2l
from mxnet import autograd, nd

train_iter, test_iter = d2l.load_data_fashion_mnist()

num_inputs = 784
num_outputs = 10

w = nd.random.normal(scale = 0.01, shape = (num_inputs, num_outputs))
b = nd.zeros(num_outputs)

params = [w, b]
for param in params:
    param.attach_grad()

def softmax(x):
    exp = x.exp()
    partition = exp.sum(axis = 1, keepdims = True)
    return exp / partition

def net(x):
    return softmax(nd.dot(x.reshape((-1, num_inputs)), w) + b)

def cross_entropy(y_hat, y):
    return -nd.pick(y_hat, y).log()

d2l.train_ch3(net, train_iter, test_iter, cross_entropy, 5, 256, params, 0.1)

d2l.show_net_acc(test_iter, net, 10)
